package control

import (
	"bytes"
	"encoding/base64"
	"fmt"
	hproxy "git.oschina.net/ant180/higgs/proxy"
	htask "git.oschina.net/ant180/higgs/task"
	"git.oschina.net/ant180/scan/config"
	"git.oschina.net/ant180/scan/model"
	ofst "git.oschina.net/ant180/util/statistic"
	"github.com/xlvector/dlog"
	"net/http"
	"sync"
)

const (
	//base64Table = "123QRSTUabcdVWXYZHijKLAWDCABDstEFGuvwxyzGHIJklmnopqr234560178912"
	base64Table = "_BCDEFGHIJpmMN8PsRSTUV1XY-9bcdefghijklLnoKqrQtuv3xyz0W2w4567OaAZ"
)

type Worker struct {
	model *model.Model
	conf  *config.Config

	msgch  chan *Message
	taskch chan string
	pxch   chan string
	wg     sync.WaitGroup

	isLocal bool
	stats   *Statistics

	reciver *Reciver
	exitch  chan bool

	ofStat ofst.Statisticer
}

func NewWorker(conf *config.Config) (*Worker, error) {
	m, err := model.NewModel(&conf.Model)
	if err != nil {
		dlog.Error("new model fail! %v", err)
		return nil, err
	}

	r, err := NewReciver(&conf.Worker.Reciver)
	if err != nil {
		dlog.Error("new reciver fail! %v", err)
		return nil, err
	}

	w := &Worker{
		conf:    conf,
		taskch:  make(chan string, conf.Worker.TaskChLen),
		msgch:   make(chan *Message, conf.Worker.MsgChLen),
		model:   m,
		isLocal: false,
		stats:   NewStatistics(),
		reciver: r,
	}

	o := ofst.NewOpenfalconOption(StatMax, StatStrs, "scan-checker")
	o.Intv = 60
	ofst, err := ofst.NewStatOF(o)
	if err != nil {
		dlog.Warn("ofstat create fail! %v", err)
	}
	w.ofStat = ofst

	if conf.Worker.UseType == "local" {
		w.isLocal = true
		w.pxch = make(chan string, conf.Worker.SPX.PXChLen)
	}

	return w, nil
}

func (p *Worker) StartAndServe() {
	p.Start()
	err := http.ListenAndServe(p.conf.Worker.ListemAddr, p)
	if err != nil {
		dlog.Fatal("listen fail: %v", err)
	}
}

func (p *Worker) Start() {
	for i := 0; i < p.conf.Worker.Concurrent; i++ {
		p.wg.Add(1)
		go p.work(i)
	}
	for i := 0; i < p.conf.Worker.MsgConcurrent; i++ {
		p.wg.Add(1)
		go p.workMsg(i)
	}

	if p.isLocal {
		for i := 0; i < p.conf.Worker.SPX.PXConCurrent; i++ {
			//p.wg.Add(1)
			go p.persistPX(i)
		}
	}

	p.ReciverStart()
}

func (p *Worker) Stop() {
	close(p.exitch)
	p.reciver.Stop()

	close(p.msgch)
	close(p.taskch)

	p.wg.Wait()
	if p.isLocal {
		close(p.pxch)
	}

	p.model.Stop()

}

var headers = map[string]string{
	//"Accept":          "application/json, text/plain, */*",
	"Accept":          "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
	"Accept-Encoding": "gzip, deflate, sdch",
	"Accept-Language": "zh - CN, zh; q = 0.8",
	"Connection":      "keep - alive",
	//"content-type":    "application/x-www-form-urlencoded; charset=utf-8",
	"User-Agent": "Mozilla/5.0 (Linux; Android 6.0; Nexus 5 Build/MRA58N) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.133 Mobile Safari/537.36",
}

func (p *Worker) workLocal(idx int) {
	downloader := htask.NewDownloaderWithTimeout(p.conf.Worker.ClientTimeout)
	proxy := &hproxy.Proxy{}

	allkeys := make([][]byte, len(p.conf.Worker.IncludeKeys), len(p.conf.Worker.IncludeKeys))
	for i, key := range p.conf.Worker.IncludeKeys {
		allkeys[i] = []byte(key)
	}

	httpApi := p.conf.Worker.HttpApi
	httpsApi := p.conf.Worker.HttpsApi
	socks4Api := p.conf.Worker.Socks4Api
	socks5Api := p.conf.Worker.Socks5Api
	var link string

	for {
		task, ok := <-p.taskch
		if !ok {
			dlog.Warn("taskch exit! %d", idx)
			return
		}
		err := proxy.Update(task)
		if err != nil {
			dlog.Warn("proxy fail! %v", err)
			continue
		}

		downloader.SetProxy(proxy)
		var body []byte
		switch proxy.Type {
		case "http":
			link = httpApi
		case "https":
			link = httpsApi
		case "socks4":
			link = socks4Api
		case "socks5":
			link = socks5Api
		default:
			continue
		}
		body, err = downloader.Get(link, headers)
		if err != nil {
			continue
		}

		for _, key := range allkeys {
			if bytes.Contains(body, key) {
				p.pxch <- task
				dlog.Info("scan found: %s", task)
				break
			}
		}
	}
}

func (p *Worker) WorkServer(idx int) {
	downloader := htask.NewDownloaderWithTimeout(p.conf.Worker.ClientTimeout)
	proxy := &hproxy.Proxy{}

	httpApi := p.conf.Worker.HttpApi
	httpsApi := p.conf.Worker.HttpsApi
	socks4Api := p.conf.Worker.Socks4Api
	socks5Api := p.conf.Worker.Socks5Api
	var link string
	var coder = base64.NewEncoding(base64Table)

	for {
		task, ok := <-p.taskch
		if !ok {
			dlog.Warn("taskch exit! %d", idx)
			return
		}
		err := proxy.Update(task)
		if err != nil {
			dlog.Warn("proxy fail! %v", err)
			continue
		}

		switch proxy.Type {
		case "http":
			link = httpApi
		case "https":
			link = httpsApi
		case "socks4":
			link = socks4Api
		case "socks5":
			link = socks5Api
		default:
			continue
		}
		downloader.SetProxy(proxy)
		resp, _ := downloader.OnlyGet(link+coder.EncodeToString([]byte(task)), headers)
		if resp != nil && resp.Body != nil {
			resp.Body.Close()
		}
	}
}

func (p *Worker) work(idx int) {
	defer p.wg.Done()

	if p.isLocal {
		p.workLocal(idx)
	} else {
		p.WorkServer(idx)
	}
}

func (p *Worker) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	for i := 0; i < SCAN_MAX_NUM; i++ {
		w.Write([]byte(fmt.Sprintf("%d: %d", i, p.stats.Get(i))))
	}
}

func (p *Worker) statInc(t ofst.StatType) {
	p.ofStat.Incr(t)
}
